"""A simple agent implementation.

DemoAgent class inherits from BaseAgent and implement all three
abstract methods: setup(), step() and reset().
"""
import json
import os
import random
from Layer import utils

from .base_agent import BaseAgent
from .map import Map


class Agent(BaseAgent):
    def __init__(self):
        self.scenario = None
        self.color = None
        self.priority = None
        self.observation = None
        self.map = None
        self.scenario_info = None
        self.map_data = None
        self.seat = None
        self.faction = None
        self.role = None
        self.controllable_ops = None
        self.team_info = None
        self.my_direction = None
        self.my_mission = None
        self.user_name = None
        self.user_id = None
        self.history = None

    def setup(self, setup_info):
        # 利用传来的推演开局相关信息进行初始化工作
        self.scenario = setup_info["scenario"]  # 想定数据
        # self.get_scenario_info(setup_info["scenario"])
        self.color = setup_info["faction"]  # faction为阵营id，0为红，1为蓝。ps：依照参考文档，faction为AI的阵营，role为AI的角色
        self.faction = setup_info["faction"]
        self.seat = setup_info["seat"]  # 本场对局的席位数，用来区分智能体的唯一值
        self.role = setup_info["role"]  # 角色，0为连长，1为营长
        self.user_name = setup_info["user_name"]
        self.user_id = setup_info["user_id"]
        # 2、创建一个行动类型为key，处理函数为value的字典。由于Python 3.6之后dict是有序的，所以字典内元素的先后顺序就暗含了不同动作之间的优先级关系
        self.priority = {
            ActionType.Occupy: self.gen_occupy,
            ActionType.Shoot: self.gen_shoot,
            ActionType.GuideShoot: self.gen_guide_shoot,
            ActionType.JMPlan: self.gen_jm_plan,
            ActionType.LayMine: self.gen_lay_mine,
            ActionType.ActivateRadar: self.gen_activate_radar,
            ActionType.ChangeAltitude: self.gen_change_altitude,
            ActionType.GetOn: self.gen_get_on,
            ActionType.GetOff: self.gen_get_off,
            # ActionType.Fork: self.gen_fork,# TODO 聚合和解聚动作会增加或减少现有所有算子id，编码起来比较复杂，暂时取消这两种动作的考虑
            # ActionType.Union: self.gen_union,
            # ActionType.EnterFort: self.gen_enter_fort,
            # ActionType.ExitFort: self.gen_exit_fort,
            ActionType.Move: self.gen_move,
            ActionType.RemoveKeep: self.gen_remove_keep,
            ActionType.ChangeState: self.gen_change_state,
            ActionType.StopMove: self.gen_stop_move,
            ActionType.WeaponLock: self.gen_WeaponLock,
            ActionType.WeaponUnFold: self.gen_WeaponUnFold,
            ActionType.CancelJMPlan: self.gen_cancel_JM_plan,
        }  # choose action by priority
        self.observation = None
        self.map = Map(
            setup_info["basic_data"],  # 地图基础数据
            setup_info["cost_data"],  # 地图通行成本数据
            setup_info["see_data"]  # 地图通视数据
        )  # use 'Map' class as a tool
        # 返回地图数据，包含坐标、地形、地势、相邻六角格坐标、与相邻六角格相连的道路类型、与相邻六角格间是否有河流。
        self.map_data = self.map.get_map_data()

    def reset(self):
        # 完成属性清空工作
        self.scenario = None
        self.color = None
        self.priority = None
        self.observation = None
        self.map = None
        self.scenario_info = None
        self.map_data = None

    def step(self, observation: dict):
        # State一般代表环境当前的所有状态。Observation一般情况下代表对于某个智能体可观测的态势。Observation是State的子集。
        # TrainEnv的step函数返回的state，表示当前环境的所有状态合集。状态合集有红方蓝方绿方态势组成：state[0]代表的是红方态势，state[1]代表的是蓝方态势,state[-1]代表的是绿方态势。
        # AI代码的step函数接受的参数是就是态势observation，它封装了当前时间，此AI(当前语境下，指的是我方agent实例)能观测到的所有盘面信息，包括算子信息、裁决信息等。
        self.observation = observation  # save observation for later use
        self.team_info = observation["role_and_grouping_info"]  # 获取玩家信息和编组信息
        self.controllable_ops = observation["role_and_grouping_info"][self.seat][
            "operators"
        ]  # 获取我方(当前语境下，指的是我方agent实例)所有算子的id
        communications = observation["communication"]  # 获取我方通信相关信息，包含了本方阵营当前的战斗任务信息等（当前语境下，指的是我方下发的所有指令）
        for command in communications:  # 遍历所有指令，如果我方我这个角色有作战任务、作战方向相关的指令，就将其保存下来
            if command["type"] in [200, 201] and command["info"]["company_id"] == self.seat:
                if command["type"] == 200:
                    self.my_mission = command
                elif command["type"] == 201:
                    self.my_direction = command
        total_actions = []

        if observation["time"]["stage"] == 1:  # 如果当前我方处于部署状态
            actions = []
            for item in observation["operators"]:  # 遍历我方所有算子
                if item["obj_id"] in self.controllable_ops:  # 如果该算子能够正常使用
                    operator = item
                    if operator["sub_type"] == 2 or operator["sub_type"] == 4:  # 如果是人或无人战车
                        actions.append(  # 让当前算子坐上初始时设定的对应的车（ "launcher": "int, 算子下车/发射后,记录所属发射器"）
                            {
                                "actor": self.seat,
                                "obj_id": operator["obj_id"],
                                "type": 303,
                                "target_obj_id": operator["launcher"],
                            }
                        )
            actions.append({  # 我方结束部署
                "actor": self.seat,
                "type": 333
            })
            return actions  # 部署阶段只做上车动作，特殊处理完毕，返回动作列表

        # 循环处理所有算子及他们的有效动作【本demo代码的核心逻辑！】
        for obj_id, valid_actions in observation["valid_actions"].items():
            if obj_id not in self.controllable_ops:  # 如果当前遍历到的算子我无法控制就跳过该算子（当前语境下，只要是我遍历到的算子我都能控制）
                continue
            for (
                    action_type
            ) in self.priority:  # 遍历预定义的优先级动作字典的key

                # TODO 如果被病毒植入：就将攻击动作、占领动作剔除
                ops_status = utils.get_ops_status(self.color)
                # if obj_id in ops_status and ops_status[obj_id]["has_implanted_virus"] == True:
                if ops_status[obj_id]["has_implanted_virus"] == True:
                    if action_type in (ActionType.Occupy, ActionType.Shoot, ActionType.GuideShoot):
                        continue

                # TODO 如果通信瘫痪：仅允许直瞄攻击动作
                network_status = utils.get_network_status(self.color)
                if random.uniform(0.00, 1.00) <= network_status[obj_id]:
                    if ActionType.Shoot not in valid_actions:  # 如果本step直瞄攻击动作不合法，则本step无动作
                        break
                    else:  # 否则就只允许进行直瞄攻击动作
                        action = self.gen_shoot(obj_id,
                                                valid_actions[ActionType.Shoot])  # 给该函数传递算子id、该动作对应的其他参数，函数返回一个参数填充完毕的动作
                        if action:  # 将填充完毕的动作加入到参数列表中
                            total_actions.append(action)
                        break  # one action per bop at a time

                if action_type not in valid_actions:  # 如果当前遍历到的动作类型非法，就跳过
                    continue
                # find the action generation method based on type
                gen_action = self.priority[action_type]  # 如果当前遍历到的动作类型合法，就拿到对应的处理函数
                action = gen_action(obj_id, valid_actions[action_type])  # 给该函数传递算子id、该动作对应的其他参数，函数返回一个参数填充完毕的动作
                if action:  # 将填充完毕的动作加入到参数列表中
                    total_actions.append(action)
                    break  # one action per bop at a time
        # if total_actions:
        #     print(
        #         f'{self.color} actions at step: {observation["time"]["cur_step"]}', end='\n\t')
        #     print(total_actions)

        return total_actions

    def get_scenario_info(self, scenario: int):  # 获取当前想定的方法
        SCENARIO_INFO_PATH = os.path.join(
            os.path.dirname(__file__), f"scenario_{scenario}.json"  # 得到想定json文件所在的路径
        )
        with open(SCENARIO_INFO_PATH, encoding="utf8") as f:
            self.scenario_info = json.load(f)  # 将想定文件内容加载到类实例的scenario_info 属性中

    def get_bop(self, obj_id):  # 根据传入的算子ID得到对应的算子对象
        """Get bop in my observation based on its id."""
        for bop in self.observation["operators"]:
            if obj_id == bop["obj_id"]:
                return bop

    def gen_occupy(self, obj_id, candidate):  # 产生占领动作
        """Generate occupy action."""
        return {
            "actor": self.seat,  # int 动作发出者席位
            "obj_id": obj_id,  # 发出动作的算子的ID int
            "type": ActionType.Occupy,  # 动作类型
        }

    def gen_shoot(self, obj_id, candidate):  # 产生具有最高攻击级别的射击动作
        """Generate shoot action with the highest attack level."""
        # candidate是一个包含候选动作的列表
        best = max(candidate, key=lambda x: x["attack_level"])  # 根据候选动作列表中每个动作的 "attack_level" 属性找到具有最高攻击级别的动作
        # 生成射击动作
        return {
            "actor": self.seat,
            "obj_id": obj_id,
            "type": ActionType.Shoot,
            "target_obj_id": best["target_obj_id"],  # 目标算子id
            "weapon_id": best["weapon_id"],  # 武器id
        }

    def gen_guide_shoot(self, obj_id, candidate):  # 产生具有最高攻击级别的引导射击动作
        """Generate guide shoot action with the highest attack level."""
        best = max(candidate, key=lambda x: x["attack_level"])
        return {
            "actor": self.seat,
            "obj_id": obj_id,  # 引导算子ID
            "type": ActionType.GuideShoot,
            "target_obj_id": best["target_obj_id"],
            "weapon_id": best["weapon_id"],
            "guided_obj_id": best["guided_obj_id"],  # 射击算子ID
        }

    def gen_jm_plan(self, obj_id, candidate):  # 间瞄射击
        """Generate jm plan action aimed at a random city."""
        weapon_id = random.choice(candidate)["weapon_id"]  # 从候选动作列表中随机选择一个动作,并提取该动作的武器id属性。
        jm_pos = random.choice([city["coord"] for city in self.observation["cities"]])  # 从观察到的数据城市列表中随机选择一个城市，并提取该城市的坐标
        return {
            "actor": self.seat,
            "obj_id": obj_id,  # 攻击算子ID
            "type": ActionType.JMPlan,
            "jm_pos": jm_pos,  # 目标位置
            "weapon_id": weapon_id,
        }

    def gen_get_on(self, obj_id, candidate):  # 按概率生成上车动作
        """Generate get on action with some probability."""
        get_on_prob = 0.5
        if random.random() < get_on_prob:  # 随机生成的数小于0.5
            target_obj_id = random.choice(candidate)["target_obj_id"]  # 从候选动作列表中随机选择一个动作,并提取该动作的成员算子id。
            return {
                "actor": self.seat,  # 动作发出者席位
                "obj_id": obj_id,  # 车辆算子ID
                "type": ActionType.GetOn,  # 动作类型：上车
                "target_obj_id": target_obj_id,  # 乘员算子ID
            }

    def gen_get_off(self, obj_id, candidate):  # 车辆算子距离随机城市距离小于10格时进行下车动作
        """Generate get off action only if the bop is within some distance of a random city."""
        bop = self.get_bop(obj_id)  # 获得车辆算子对象
        destination = random.choice(
            [city["coord"] for city in self.observation["cities"]]
        )  # 从观察到的城市中任选一个，得到其坐标
        if bop and self.map.get_distance(bop["cur_hex"], destination) <= 10:
            # 车辆仍然存在并且车辆的当前位置到随机选择的城市的距离小于或等于10
            target_obj_id = random.choice(candidate)["target_obj_id"]  # 从候选动作列表中随机选择一个动作，并提取该动作的目标算子id属性。
            return {
                "actor": self.seat,
                "obj_id": obj_id,  # 车辆算子ID
                "type": ActionType.GetOff,
                "target_obj_id": target_obj_id,  # 乘员算子ID
            }

    def gen_change_state(self, obj_id, candidate):  # 以一定概率生成改变状态动作
        """Generate change state action with some probability."""
        change_state_prob = 0.001
        if random.random() < change_state_prob:  # 改变状态的概率为0.001
            target_state = random.choice(candidate)["target_state"]  # 候选动作列表中随机选择一个动作，并提取该动作的目标状态属性。
            return {
                "actor": self.seat,
                "obj_id": obj_id,  # 算子ID
                "type": ActionType.ChangeState,
                "target_state": target_state,  # 目标状态 0-正常机动 1-行军 2-一级冲锋 3-二级冲锋, 4-掩蔽 5-半速x
            }

    def gen_remove_keep(self, obj_id, candidate):  # 以一定概率生成移除压制动作
        """Generate remove keep action with some probability."""
        remove_keep_prob = 0.2
        if random.random() < remove_keep_prob:
            return {
                "actor": self.seat,  # 动作发出者席位
                "obj_id": obj_id,  # 算子ID int
                "type": ActionType.RemoveKeep,
            }

    def gen_move(self, obj_id, candidate):
        """Generate move action to a random city."""
        bop = self.get_bop(obj_id)
        if bop["sub_type"] == 3:
            return
        destination = random.choice(
            [city["coord"] for city in self.observation["cities"]]
        )
        if self.my_direction:
            destination = self.my_direction["info"]["target_pos"]
        if bop and bop["cur_hex"] != destination:
            move_type = self.get_move_type(bop)
            route = self.map.gen_move_route(bop["cur_hex"], destination, move_type)
            return {
                "actor": self.seat,
                "obj_id": obj_id,
                "type": ActionType.Move,
                "move_path": route,
            }

    def get_move_type(self, bop):
        """Get appropriate move type for a bop."""
        bop_type = bop["type"]
        if bop_type == BopType.Vehicle:
            if bop["move_state"] == MoveType.March:
                move_type = MoveType.March
            else:
                move_type = MoveType.Maneuver
        elif bop_type == BopType.Infantry:
            move_type = MoveType.Walk
        else:
            move_type = MoveType.Fly
        return move_type

    def gen_stop_move(self, obj_id, candidate):
        """Generate stop move action only if the bop is within some distance of a random city.

        High probability for the bop with passengers and low for others.
        """
        bop = self.get_bop(obj_id)
        destination = random.choice(
            [city["coord"] for city in self.observation["cities"]]
        )
        if self.map.get_distance(bop["cur_hex"], destination) <= 10:
            stop_move_prob = 0.9 if bop["passenger_ids"] else 0.01
            if random.random() < stop_move_prob:
                return {
                    "actor": self.seat,
                    "obj_id": obj_id,
                    "type": ActionType.StopMove,
                }

    def gen_WeaponLock(self, obj_id, candidate):
        bop = self.get_bop(obj_id)
        prob_weaponlock = 0.001
        if (
                max(self.map_data[bop["cur_hex"] // 100][bop["cur_hex"] % 100]["roads"]) > 0
                or random.random() < prob_weaponlock
        ):
            return {"actor": self.seat, "obj_id": obj_id, "type": ActionType.WeaponLock}

    def gen_WeaponUnFold(self, obj_id, candidate):
        bop = self.get_bop(obj_id)
        destination = random.choice(
            [city["coord"] for city in self.observation["cities"]]
        )
        if self.map.get_distance(bop["cur_hex"], destination) <= 10:
            return {
                "actor": self.seat,
                "obj_id": obj_id,
                "type": ActionType.WeaponUnFold,
            }

    def gen_cancel_JM_plan(self, obj_id, candidate):
        cancel_prob = 0.0001
        if random.random() < cancel_prob:
            return {
                "actor": self.seat,
                "obj_id": obj_id,
                "type": ActionType.CancelJMPlan,
            }

    def gen_grouping_info(self, observation):
        def partition(lst, n):
            return [lst[i::n] for i in range(n)]

        operator_ids = []
        for operator in observation["operators"] + observation["passengers"]:
            if operator["color"] == self.color:
                operator_ids.append(operator["obj_id"])
        lists_of_ops = partition(operator_ids, len(self.team_info.keys()))
        grouping_info = {"actor": self.seat, "type": 100}
        info = {}
        for teammate_id in self.team_info.keys():
            info[teammate_id] = {"operators": lists_of_ops.pop()}
        grouping_info["info"] = info
        return [grouping_info]

    def gen_battle_direction_info(self, observation):
        direction_info = []
        for teammate_id in self.team_info.keys():
            direction = {
                "actor": self.seat,
                "type": 201,
                "info": {
                    "company_id": teammate_id,
                    "target_pos": random.choice(observation["cities"])["coord"],
                    "start_time": 0,
                    "end_time": 1800,
                },
            }
            direction_info.append(direction)
        return direction_info

    def gen_battle_mission_info(self, observation):
        mission_info = []
        for teammate_id in self.team_info.keys():
            mission = {
                "actor": self.seat,
                "type": 200,
                "info": {
                    "company_id": teammate_id,
                    "mission_type": random.randint(0, 2),
                    "target_pos": random.choice(observation["cities"])["coord"],
                    "route": [
                        random.randint(0, 9000),
                        random.randint(0, 9000),
                        random.randint(0, 9000),
                    ],
                    "start_time": 0,
                    "end_time": 1800,
                },
            }
            mission_info.append(mission)
        return mission_info

    def gen_fork(self, obj_id, candidate):
        prob = 0.01
        if random.random() < prob:
            return None
        return {"actor": self.seat, "obj_id": obj_id, "type": ActionType.Fork}

    def gen_union(self, obj_id, candidate):
        prob = 0.1
        if random.random() < prob:
            return None
        return {
            "actor": self.seat,
            "obj_id": obj_id,
            "target_obj_id": random.choice(candidate)["target_obj_id"],
            "type": ActionType.Union,
        }

    def gen_change_altitude(self, obj_id, candidate):
        prob = 0.05
        if random.random() < prob:
            return None
        return {
            "actor": self.seat,
            "obj_id": obj_id,
            "target_obj_id": random.choice(candidate)["target_altitude"],
            "type": ActionType.ChangeAltitude,
        }

    def gen_activate_radar(self, obj_id, candidate):
        prob = 1
        if random.random() < prob:
            return None
        return {"actor": self.seat, "obj_id": obj_id, "type": ActionType.ActivateRadar}

    def gen_enter_fort(self, obj_id, candidate):
        prob = 0.5
        if random.random() < prob:
            return None
        return {
            "actor": self.seat,
            "obj_id": obj_id,
            "type": ActionType.EnterFort,
            "target_obj_id": random.choice(candidate)["target_obj_id"],
        }

    def gen_exit_fort(self, obj_id, candidate):
        prob = 0.1
        if random.random() < prob:
            return None
        return {
            "actor": self.seat,
            "obj_id": obj_id,
            "type": ActionType.ExitFort,
            "target_obj_id": random.choice(candidate)["target_obj_id"],
        }

    def gen_lay_mine(self, obj_id, candidate):
        prob = 1
        if random.random() < prob:
            return None
        return {
            "actor": self.seat,
            "obj_id": obj_id,
            "type": 20,
            "target_pos": random.randint(0, 9177),
        }


class BopType:
    Infantry, Vehicle, Aircraft = range(1, 4)


class ActionType:
    (
        Move,
        Shoot,
        GetOn,
        GetOff,
        Occupy,
        ChangeState,
        RemoveKeep,
        JMPlan,
        GuideShoot,
        StopMove,
        WeaponLock,
        WeaponUnFold,
        CancelJMPlan,
        Fork,
        Union,
        ChangeAltitude,
        ActivateRadar,
        EnterFort,
        ExitFort,
        LayMine,
    ) = range(1, 21)


class MoveType:
    Maneuver, March, Walk, Fly = range(4)
